from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Table
from sqlalchemy.orm.session import Session

from airflow.models.base import Base
from airflow.utils.session import provide_session


class ProjectDatasource(Base):
    """the class to get user info."""

    __tablename__ = "idata_project_datasource"
    __table_args__ = {"useexisting": True}

    id = Column(Integer, primary_key=True, autoincrement=True)
    project_id = Column(Integer)
    datasource_id = Column(Integer)

    @staticmethod
    @provide_session
    def get_project_datasource(id, session: Session = None):
        return session.query(ProjectDatasource).filter(ProjectDatasource.id == id).first()

    @staticmethod
    @provide_session
    def add_project_datasource(project_datasource, session: Session = None):
        result = session.query(ProjectDatasource)\
            .filter(ProjectDatasource.project_id == project_datasource.project_id, ProjectDatasource.datasource_id == project_datasource.datasource_id)\
            .first()
        if result is None:
            return session.add(project_datasource)

    @staticmethod
    @provide_session
    def delete_project_datasource(project_id, datasource_id, session: Session = None):
        query = session.query(ProjectDatasource)
        if project_id != '' and project_id:
            query = query.filter(ProjectDatasource.project_id == project_id)
        if datasource_id != '' and datasource_id:
            query = query.filter(ProjectDatasource.datasource_id == datasource_id)
        query.delete()

    @staticmethod
    @provide_session
    def update_project_datasource(project_datasource, session: Session = None):
        result = session.query(ProjectDatasource).filter(ProjectDatasource.id == project_datasource.id).first()
        if result is None:
            return
        else:
            return session.merge(project_datasource)

    def to_json(self):
        if self is None:
            return None
        data = {
            'id': self.id,
            'project_id': self.project_id,
            'datasource_id': self.datasource_id
        }
        return data
